#include "KDTree.h"
#include "Object.h"
#include "AABB.h"
#include <algorithm>
#include <cmath>
#include <stack>
#include <queue>
#include <functional>
#include "Ray.h"

struct mray::KDNode{

	enum TYPE : unsigned char{ INTER, LEAF, UNKNOWN, };

	void initLeaf(const std::vector<const AABB*>& boxes);
	void initInter(KDNode* lchild, KDNode* rchild);

	bool isLeaf()const{return m_type == LEAF;}
	bool isInter()const {return m_type == INTER;}

	~KDNode();

	union{
		struct{
			KDNode* m_lchild;
			KDNode* m_rchild;
		};
		struct{
			const Collible **	m_collibles;
			size_t		m_num;
		};
	};

	TYPE	m_type;
	float	m_split;
	AXIS	m_axis;
};

void mray::KDNode::initLeaf(const std::vector<const AABB*>& boxes)
{
	m_type = LEAF;
	m_collibles = new const Collible*[boxes.size()];
	m_num = boxes.size();
	for(size_t i = 0; i < boxes.size(); i++){
		m_collibles[i] = boxes[i]->m_object;
	}
}

mray::KDNode::~KDNode()
{
	switch(m_type){
	case INTER:
		delete m_lchild;
		delete m_rchild;
		break;
	case LEAF:
		delete [] m_collibles;
	default:
		break;
	}
}

void mray::KDNode::initInter(KDNode* lchild, KDNode* rchild)
{
	m_type = INTER;
	m_lchild = lchild;
	m_rchild = rchild;
}

void mray::KDTree::Build(const std::vector<Collible*>& collibles)
{
	//create all AABBs
	m_tmpBoxes.resize(collibles.size());
	for(int i = 0; i < collibles.size(); i++){
		collibles[i]->GetAABB(m_tmpBoxes[i]);
		m_tmpBoxes[i].m_object = collibles[i];
		m_sceneBox &= m_tmpBoxes[i];
	}

	std::vector<const AABB*> boxes;
	boxes.resize(m_tmpBoxes.size());
	for(size_t i = 0; i < m_tmpBoxes.size(); i++)
		boxes[i] = &(m_tmpBoxes[i]);

	m_root = new KDNode;
	m_maxDepth = int(std::log(collibles.size())*1.3)+8;

	RecursivelyBuild(m_sceneBox, boxes, m_root, 0);
}

struct mray::KDTODO{
	KDTODO(KDNode* node, float min, float max){
		m_node = node;
		m_min = min;
		m_max = max;
	}
	KDNode* m_node;
	float m_max;
	float m_min;
};

bool mray::KDTree::Hit(Ray& ray, Result* pResult /*= nullptr*/) const 
{
	vec3 invDir(1.f/ray.m_dir.x,1.f/ray.m_dir.y,1.f/ray.m_dir.z);

	std::stack<KDTODO> todoStack;

	//ray and AABB test
	bool ret = m_sceneBox.Hit(ray);
	if(ret == false) return false;
	
	//for compatibility
	if(ray.m_onlyMin)
		todoStack.push(KDTODO(m_root, ray.m_maxt, ray.m_mint));
	else
		todoStack.push(KDTODO(m_root, ray.m_mint, ray.m_maxt));

	ray.m_mint = std::numeric_limits<float>::max();

	while(todoStack.empty() == false){

		KDNode* node = todoStack.top().m_node;
		float tmin = todoStack.top().m_min;
		float tmax = todoStack.top().m_max;

		if(node->m_type == KDNode::INTER){
			//process inter kdnode
			AXIS axis = node->m_axis;
			KDNode* fchild, *schild;

			float tplane = (node->m_split - ray.m_ori[axis])*invDir[axis];

			bool belowFirst = (ray.m_ori[axis] < node->m_split || (ray.m_ori[axis] == node->m_split && ray.m_dir[axis] >= 0.f));
			if(belowFirst){
				fchild = node->m_lchild;
				schild = node->m_rchild;
			}else{
				fchild = node->m_rchild;
				schild = node->m_lchild;
			}

			if(tplane > tmax || tplane <= 0.f){
				//only process first node
				todoStack.top().m_node = fchild;
			}else if(tplane < tmin){
				//only process second node
				todoStack.top().m_node = schild;
			}else{
				//process both above and below
				todoStack.top().m_node = schild;
				todoStack.top().m_min = tplane;
				todoStack.top().m_max = tmax;
				todoStack.push(KDTODO(fchild, tmin, tplane));
			}
		}else{
			//process leaf kdnode
			ret = false;
			ray.m_mint = std::numeric_limits<float>::max();

			for(size_t i = 0; i < node->m_num; i++)
				ret |= node->m_collibles[i]->Hit(ray, pResult);

			if(ret && ray.m_mint < tmax) break;

			if(todoStack.empty() == false)
				todoStack.pop();
		}

	}

	return ret;
}

void mray::KDTree::RecursivelyBuild(const AABB& nodeBox, const std::vector<const AABB*>& boxes, KDNode* node, int depth)
{

	if(boxes.size() <= 2 || depth > m_maxDepth){
		node->initLeaf(boxes);
		return;
	}

	//prepare edges
	for(uchar axis = 0; axis < 3; axis++){
		auto& edges = m_tmpEdges[axis];
		edges.resize(boxes.size()*2);
		for(size_t i = 0; i < boxes.size(); i++){
			boxes[i]->GetEdge(edges[2*i], (AXIS)axis, Edge::START);
			boxes[i]->GetEdge(edges[2*i+1], (AXIS)axis, Edge::END);
		}
		std::sort(edges.begin(), edges.end());
	}

	//choose the best edge
	float bestSplit;
	float bestCost = FLT_MAX;
	size_t bestOffset;
	AXIS bestAxis = AXIS::NUM;

	float totalSA = nodeBox.GetSurfaceArea();
	float invTotalSA = 1.f/totalSA;

	for(uchar axis = (uchar)nodeBox.GetMaxExtent(), nTry = 0; nTry < 3; axis = (axis+1)%3, nTry++){
		const auto& edges = m_tmpEdges[axis];
		size_t nBelow = 0, nAbove = boxes.size();
		for(size_t i = 0; i < edges.size(); i++){
			if(edges[i].m_type == Edge::END) nAbove--;
			float edget = edges[i].m_t;
			if(edget > nodeBox.m_min[axis] && edget < nodeBox.m_max[axis]){
				float belowSA = nodeBox.GetSurfaceArea((AXIS)axis, edget, true);
				float aboveSA = nodeBox.GetSurfaceArea((AXIS)axis, edget, false);
				float pBelow = belowSA * invTotalSA;
				float pAbove = aboveSA * invTotalSA;
				float eb = (nAbove == 0 || nBelow == 0) ? m_emptyBonus : 0.f;
				float cost = (1.f-eb)*(pBelow*nBelow + pAbove*nAbove);
				if(cost < bestCost){
					bestSplit = edges[i].m_t;
					bestCost = cost;
					bestOffset = i;
					bestAxis = (AXIS)axis;
				}
			}
			if(edges[i].m_type == Edge::START) nBelow++;
		}
	}
	
	//if not found any axis
	if(bestAxis == AXIS::NUM){
		node->initLeaf(boxes);
		return;
	}

	//need to handle the bad split
	float start = nodeBox.m_min[bestAxis];
	float end   = nodeBox.m_max[bestAxis];
	float len   = end - start;
	float globalLen = m_sceneBox.m_max[bestAxis] - m_sceneBox.m_min[bestAxis];
	if(len/globalLen < m_edgeProtection){
		node->initLeaf(boxes);
		return;
	}
	//not a good spli
	if((end-bestSplit)/len < m_edgeProtection || (bestSplit-start)/len < m_edgeProtection)
	{
		node->initLeaf(boxes);
		return;
	}

	node->m_axis = bestAxis;
	node->m_split = bestSplit;

	//split the boxes to the next routine
	std::vector<const AABB*> boxes0;
	std::vector<const AABB*> boxes1;
	const auto& edges = m_tmpEdges[bestAxis];
	for(size_t i = 0; i < bestOffset; i++){
		const auto& edge = edges[i];
		if(edge.m_type == Edge::START)
			boxes0.push_back(edge.m_aabb);
	}
	for(size_t i = bestOffset+1; i < boxes.size()*2; i++){
		const auto& edge = edges[i];
		if(edge.m_type == Edge::END)
			boxes1.push_back(edge.m_aabb);
	}
	AABB belowBox, aboveBox;
	nodeBox.Split((AXIS)bestAxis, bestSplit, belowBox, aboveBox);
	//recursively build
	node->initInter(new KDNode,  new KDNode);

	RecursivelyBuild(belowBox, boxes0, node->m_lchild, depth+1);
	RecursivelyBuild(aboveBox, boxes1, node->m_rchild, depth+1);
}

mray::KDTree::KDTree()
{
	m_emptyBonus = 0.f;
	m_edgeProtection = 0.05f;
	m_ofsbin.open("kdlog.bin", std::ios::out | std::ios::binary);
	m_ofstxt.open("kdlog.txt", std::ios::out );
}

mray::KDTree::~KDTree()
{
	delete m_root;
	m_ofsbin.close();
}

void mray::KDTree::log(const AABB& nodeBox)
{
	static std::array<float, 6> tmp;
	tmp[0] = nodeBox.m_min.x;
	tmp[1] = nodeBox.m_min.y;
	tmp[2] = nodeBox.m_min.z;
	tmp[3] = nodeBox.m_max.x;
	tmp[4] = nodeBox.m_max.y;
	tmp[5] = nodeBox.m_max.z;
	m_ofsbin.write((const char*)(tmp.data()), sizeof(tmp));
	m_ofsbin.flush();
	
	int i = 0;
	for(; i < 3; i++)
		m_ofstxt << tmp[i] << ' ';
	m_ofstxt << std::endl;

	for(; i < 6; i++)
		m_ofstxt << tmp[i] << ' ';
	m_ofstxt << std::endl << std::endl;
	m_ofstxt.flush();

}

void mray::KDTree::dump() const
{
//	std::ofstream ofs("kddump.txt");
	std::function<void(KDNode*, int)> traverse = [&](KDNode* node, int depth)->void{
		std::cout << "depth:" << depth << (node->isInter() ? " inte" : " leaf") << std::endl;
		if(node->isInter()){
			traverse(node->m_lchild, depth+1);
			traverse(node->m_rchild, depth+1);
		}
	};
	traverse(m_root, 0);
}

